import getopt
import sys

import pandas as pd
from elasticsearch import Elasticsearch
# from elasticsearch.exceptions import ConnectionTimeout

from air_web.web_flask.dal.base_dal import EsBaseDal
from air_web.config.config import config
from air_web.data_platform import sql_engine


def proc_agg_data(agg_res):
    res_list = []
    for bucket in agg_res:
        key_list = bucket["key"].split("|")
        bucket_dict = {"data_time": key_list[0]}
        for key, val_dict in bucket.items():
            if key in ["key", "doc_count"]:
                continue
            bucket_dict.update({key: val_dict["value"]})
        res_list.append(bucket_dict)
    df = pd.DataFrame(res_list)
    return df


def get_cons(cons_df, start_time, end_time):
    es = EsBaseDal(config["ES_HOST"])
    on7_dict = (
        cons_df[["on7", "on5"]]
        .drop_duplicates()
        .set_index("on7")["on5"]
        .to_dict()
    )
    res_df = pd.DataFrame()
    for on7, on5 in on7_dict.items():
        print(f"on7:{on7} ...")

        is_cal = "1" if config["c_cons_filter_is_cal"] else "*"
        index_name = (
            config["POWERLOAD"]
            .replace("on5", str(on5))
            .replace("on7", str(on7))
            .replace("cal01", is_cal)
        )
        rules = [("data_time", "query", '>=', start_time),
                 ("data_time", "query", '<', end_time),
                 ("type_code_sort", "query", '=', 1)]
        on7_df = es.query_dataframe(rules, index_name=index_name, doc_time_field="data_time",
                                    source=["cons_no", "data_time","p_total", "p_kt"])
        if len(on7_df) == 0:
            print(f"数据为空，on7:{on7}")
            continue
        res_df = pd.concat([res_df, on7_df])
    if res_df.empty:
        print("结果为空")
        return pd.DataFrame()
    res_df = pd.merge(cons_df, res_df, on="cons_no", how="right")
    res_df = res_df.sort_values(["on7", "cons_no", "data_time"])
    return res_df


def get_cons_info():
    sql = """select distinct cons_no, cons_name, on5, shi, on7, xian, type_code, pare_type_code
              from c_cons
              where type_code_sort=1
              order by on7"""
    cons_df = sql_engine.query(sql)

    print(f"从c_cons读取用户数：{len(cons_df)}")

    return cons_df


def main(start_time, end_time):
    cons_df = get_cons_info()
    res_df = get_cons(cons_df, start_time, end_time)
    file_name = f'cons_res.csv'
    res_df.to_csv(file_name, index=False)
    print(f"结果已生成:{file_name}")


if __name__ == '__main__':
    es = Elasticsearch(config["ES_HOST"])
    start_time = '2024-03-15'
    end_time = '2024-03-17'

    opts, args = getopt.getopt(sys.argv[1:], 's:e:')
    for opt, val in opts:
        if opt == '-s':
            start_time = val
        elif opt == '-e':
            end_time = val
    main(start_time, end_time)
